"""
数据模型模块

定义聊天助手相关的数据结构和类型。
"""

from dataclasses import dataclass
from typing import List, Literal, Optional, Dict, Any
from datetime import datetime
from enum import Enum


class MessageRole(str, Enum):
    """消息角色枚举"""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"


@dataclass
class ChatMessage:
    """
    聊天消息数据结构
    
    Attributes:
        role: 消息角色（system/user/assistant）
        content: 消息内容
        timestamp: 消息时间戳
        metadata: 附加元数据
    """
    
    role: MessageRole
    content: str
    timestamp: Optional[datetime] = None
    metadata: Optional[Dict[str, Any]] = None
    
    def __post_init__(self) -> None:
        """初始化后处理"""
        if self.timestamp is None:
            self.timestamp = datetime.now()
        
        if self.metadata is None:
            self.metadata = {}
    
    def to_langchain_format(self) -> Dict[str, str]:
        """
        转换为LangChain消息格式
        
        Returns:
            Dict[str, str]: LangChain格式的消息
        """
        return {
            "role": self.role.value,
            "content": self.content
        }
    
    def __str__(self) -> str:
        """字符串表示"""
        timestamp_str = self.timestamp.strftime("%Y-%m-%d %H:%M:%S") if self.timestamp else "Unknown"
        return f"[{timestamp_str}] {self.role.value}: {self.content}"


@dataclass
class ChatSession:
    """
    聊天会话数据结构
    
    Attributes:
        session_id: 会话ID
        messages: 消息列表
        created_at: 创建时间
        updated_at: 更新时间
        metadata: 会话元数据
    """
    
    session_id: str
    messages: List[ChatMessage]
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    metadata: Optional[Dict[str, Any]] = None
    
    def __post_init__(self) -> None:
        """初始化后处理"""
        if self.created_at is None:
            self.created_at = datetime.now()
        
        if self.updated_at is None:
            self.updated_at = self.created_at
        
        if self.metadata is None:
            self.metadata = {}
    
    def add_message(self, message: ChatMessage) -> None:
        """
        添加消息到会话
        
        Args:
            message: 要添加的消息
        """
        self.messages.append(message)
        self.updated_at = datetime.now()
    
    def get_messages_for_langchain(self) -> List[Dict[str, str]]:
        """
        获取适合LangChain的消息格式
        
        Returns:
            List[Dict[str, str]]: LangChain格式的消息列表
        """
        return [msg.to_langchain_format() for msg in self.messages]
    
    def get_recent_messages(self, count: int = 10) -> List[ChatMessage]:
        """
        获取最近的消息
        
        Args:
            count: 要获取的消息数量
            
        Returns:
            List[ChatMessage]: 最近的消息列表
        """
        return self.messages[-count:] if len(self.messages) > count else self.messages
    
    def clear_messages(self) -> None:
        """清空所有消息"""
        self.messages.clear()
        self.updated_at = datetime.now()
    
    def get_message_count(self) -> int:
        """获取消息总数"""
        return len(self.messages)
    
    def __str__(self) -> str:
        """字符串表示"""
        return f"ChatSession(id={self.session_id}, messages={len(self.messages)}, updated={self.updated_at})"


@dataclass
class ChatResponse:
    """
    聊天响应数据结构
    
    Attributes:
        content: 响应内容
        success: 是否成功
        error_message: 错误信息（如果有）
        metadata: 响应元数据
        timestamp: 响应时间戳
    """
    
    content: str
    success: bool = True
    error_message: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    timestamp: Optional[datetime] = None
    
    def __post_init__(self) -> None:
        """初始化后处理"""
        if self.timestamp is None:
            self.timestamp = datetime.now()
        
        if self.metadata is None:
            self.metadata = {}
    
    @classmethod
    def success_response(cls, content: str, metadata: Optional[Dict[str, Any]] = None) -> "ChatResponse":
        """
        创建成功响应
        
        Args:
            content: 响应内容
            metadata: 响应元数据
            
        Returns:
            ChatResponse: 成功响应实例
        """
        return cls(content=content, success=True, metadata=metadata)
    
    @classmethod
    def error_response(cls, error_message: str, metadata: Optional[Dict[str, Any]] = None) -> "ChatResponse":
        """
        创建错误响应
        
        Args:
            error_message: 错误信息
            metadata: 响应元数据
            
        Returns:
            ChatResponse: 错误响应实例
        """
        return cls(content="", success=False, error_message=error_message, metadata=metadata)
